package com.shujia.mllib

import org.apache.spark.ml.classification.{LogisticRegression, LogisticRegressionModel}
import org.apache.spark.sql.{DataFrame, Dataset, Row, SparkSession}

object Demo6TrainModel {
  def main(args: Array[String]): Unit = {
    val spark: SparkSession = SparkSession
      .builder()
      .master("local[6]")
      .appName("image")
      .getOrCreate()
    import spark.implicits._
    import org.apache.spark.sql.functions._


    val imageData: DataFrame = spark
      .read
      .format("libsvm")
      .load("data/image")


    /**
      * 将数据拆分成训练集和测试集
      *
      */

    val split: Array[Dataset[Row]] = imageData.randomSplit(Array(0.8, 0.2))
    val train: DataFrame = split(0)
    val test: DataFrame = split(1)

    /**
      * 构建算法，训练模型
      *
      */

    //逻辑回归
    val logisticRegression = new LogisticRegression()

    //将训练集带入算法，训练模型
    val model: LogisticRegressionModel = logisticRegression.fit(train)


    /**
      * 使用测试集测试模型的准确率
      *
      */

    val testDF: DataFrame = model.transform(test)


    //计算准确率
    val p: Double = testDF.where($"label" === $"prediction").count().toDouble / testDF.count()

    println(s"准确率：$p")

    //保存模型
    model.write.overwrite().save("data/image_model")

  }

}
